#!/usr/bin/python3
# -*- coding: utf-8 -*-
# Created by Ross on 19-3-26
import os
from collections import Counter

import numpy as np
import pandas as pd
import tensorflow.contrib.eager as tfe
from keras.preprocessing.sequence import pad_sequences
from setuptools.command.saveopts import saveopts

from hparams import Hparams
from model.JointModel import JointModel
from utils import load_hparams, get_checkpoints, id2label, Calculate

RESULT_FILE = 'results.csv'
hparams = Hparams()
hp = hparams.parser.parse_args()  # 超参数字典
load_hparams(hp, hp.log_dir)
hp_keys = sorted(vars(hp).keys())
hp_values = [str(vars(hp)[k]) for k in hp_keys]

test_x = np.load(os.path.join(hp.data_dir, 'test_x.npy'))
test_y = np.load(os.path.join(hp.data_dir, 'test_y.npy')).astype(np.int32)
if hp.num_class == 2:
    test_y = np.where(test_y == 3, 1, 0)
test_lens = [min(len(x), hp.seq_maxlen) for x in test_x]
test_x = pad_sequences(test_x, hp.seq_maxlen, 'float32', padding='post', truncating='post')
if hp.fake_task == 'POS':
    test_POS = np.load(os.path.join(hp.data_dir, 'POStest_x.npy'))
    test_POS = pad_sequences(test_POS, hp.seq_maxlen, 'float32', padding='post', truncating='post')

with open('data/full_data/segment/test_x.txt', 'r', encoding='utf-8') as fp:
    sentences = [line.strip() for line in fp]


def predict():
    model = JointModel(hp.seq_maxlen, hp.emb_size, hp.rnn_size, hp.rnn_keep_prob, hp.domain_keep_prob,
                       hp.use_self_att, hp.soft_att_size,
                       hp.fake_task, hp.num_class, hp.use_crf, ntags=hp.ntags)
    ckpts = get_checkpoints(hp.log_dir)
    acc = tfe.metrics.Accuracy()
    pred = []
    single_preds = []
    for ckpt in ckpts:
        acc.init_variables()
        model.load_weights(ckpt)
        domain_pred = model.predict_domain([test_x, test_lens], top_k=2, save_att=False).indices  # TOP2
        single_preds.append(domain_pred.numpy())
        pred.append(domain_pred.numpy()[:, 0])  # 把第一个拿出来去联合预测
    result = []
    for i in zip(*pred):
        c = Counter(i)
        result.append(c.most_common(1)[0][0])
    acc(test_y.flatten(), result)
    df = pd.DataFrame(list(zip(sentences, result, test_y.flatten())), columns=['sentence', 'prediction', 'ground_true'])
    df['预测'] = df['prediction'].map(lambda x: id2label[x])
    df['正确'] = df['ground_true'].map(lambda x: id2label[x])
    # 每一折top2
    for i, single_pred in enumerate(single_preds):
        df['fold_{}_top2'.format(i)] = list(map(lambda x: (id2label[x[0]], id2label[x[1]]), single_pred))

    result_path = os.path.join(hp.log_dir, '%.5f.csv' % acc.result().numpy())
    df.to_csv(result_path, index=False, encoding='utf-8')
    calculate = Calculate(result_path)
    calculate.get_confusion_matrix_31_csv(result_path.replace('.csv', '_matrix.csv'))
    print(acc.result().numpy())
    if not os.path.exists(RESULT_FILE):
        with open(RESULT_FILE, 'w', encoding='utf-8') as f:
            f.write('test_acc,')
            f.write(','.join(hp_keys))
            f.write('\n')
    with open(RESULT_FILE, 'a', encoding='utf-8') as f:
        f.write(str(acc.result().numpy()) + ',')
        f.write(','.join(hp_values))
        f.write('\n')
    print('Wrote to {}'.format(RESULT_FILE))


if __name__ == '__main__':
    predict()
